#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

from dir_info import raw_X_data_dir

data_reader = DataReader(db='validation')

## 路径不准确
hf_data_dir = './zbc_factors_lib/hf_basic/factor'

def main(start_date='2016-01-01', end_date='2020-01-31', is_update=False, is_overwrite=False):
    # TODO - 处理高频trade资金流数据
    hf_order_quantile_data = pd.read_hdf(os.path.join(hf_data_dir, 'hf_order_quantile_data.h5'))

    hf_order_quantile_data = hf_order_quantile_data.set_index(['date', 'stock_code'])

    hf_order_quantile_data = np.log(hf_order_quantile_data)
    hf_order_quantile_data.replace([np.inf, -np.inf], 0.0, inplace=True)

    hf_order_quantile_data = hf_order_quantile_data.astype(np.float32)

    hf_order_quantile_data = hf_order_quantile_data.drop(['total_money', 'total_volume'], axis=1)

    hf_order_quantile_data.columns = ['log_' +c for c in hf_order_quantile_data.columns]

    # tmp = hf_order_quantile_data[hf_order_quantile_data.index.get_level_values('stock_code') == '300104.SZ']
    # tmp.loc[tmp.index.get_level_values('date') <= pd.to_datetime('2018-01-24')]

    # TODO - 读取处理日数据
    processed_daily_stock_trade_data = data_reader.read_basic_data_table('processed_daily_stock_trade_data')

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.rename(columns={'code': 'stock_code'})

    processed_daily_stock_trade_data = processed_daily_stock_trade_data.set_index(['date', 'stock_code'])

    processed_daily_stock_trade_data['bclose'] = processed_daily_stock_trade_data['close'] * processed_daily_stock_trade_data['back_adjfactor']

    processed_daily_stock_trade_data['swing'] = processed_daily_stock_trade_data['high'] / processed_daily_stock_trade_data['low'] - 1.0

    processed_daily_stock_trade_data['oc_ret'] = processed_daily_stock_trade_data['close'] / processed_daily_stock_trade_data['open'] - 1.0

    processed_daily_stock_trade_data['log_volume'] = np.log(processed_daily_stock_trade_data['volume'])
    processed_daily_stock_trade_data['log_money'] = np.log(processed_daily_stock_trade_data['money'])
    processed_daily_stock_trade_data.loc[processed_daily_stock_trade_data['volume']==0, 'log_volume'] = np.nan
    processed_daily_stock_trade_data.loc[processed_daily_stock_trade_data['money']==0, 'log_money'] = np.nan

    ret = processed_daily_stock_trade_data.groupby('stock_code', as_index=False)['bclose'].pct_change()

    ret.index = ret.index.droplevel(level=0)

    processed_daily_stock_trade_data['ret'] = ret.loc[processed_daily_stock_trade_data.index]

    processed_daily_stock_trade_data = processed_daily_stock_trade_data[
        [
            'ret',
            'oc_ret',
            'swing',
            'log_volume',
            'log_money',
            'turnover',
        ]
    ]

    ##
    processed_daily_stock_trade_data = processed_daily_stock_trade_data.astype(np.float32)

    processed_daily_stock_trade_data = \
        processed_daily_stock_trade_data[(processed_daily_stock_trade_data.index.get_level_values('date') >= start_date) &
                                         (processed_daily_stock_trade_data.index.get_level_values('date') <= end_date)]

    # TODO - 读取处理市值，估值因子
    valuation_ep_ttm = data_reader.read_factor_table('valuation_ep_ttm')
    valuation_bp = data_reader.read_factor_table('valuation_bp')
    scale_circulate_market_size = data_reader.read_factor_table('scale_circulate_market_size')

    scale_circulate_market_size['scale_circulate_market_size'] = np.log(scale_circulate_market_size['scale_circulate_market_size'])

    valuation_ep_ttm = valuation_ep_ttm.set_index(['date', 'stock_code'])
    valuation_bp = valuation_bp.set_index(['date', 'stock_code'])
    scale_circulate_market_size = scale_circulate_market_size.set_index(['date', 'stock_code'])
    scale_circulate_market_size = scale_circulate_market_size.rename(columns={'scale_circulate_market_size': 'log_cir_cap'})

    valuation_ep_ttm = valuation_ep_ttm.astype(np.float32)
    valuation_bp = valuation_bp.astype(np.float32)
    scale_circulate_market_size = scale_circulate_market_size.astype(np.float32)

    # TODO - merge
    X_hf_order_quantile_data_v1 = pd.concat([processed_daily_stock_trade_data,
                                             hf_order_quantile_data,
                                             valuation_ep_ttm,
                                             valuation_bp,
                                             scale_circulate_market_size],
                                             axis=1,
                                             join_axes=[processed_daily_stock_trade_data.index])

    # TODO - 保存
    if not is_update:
        X_hf_order_quantile_data_v1.to_hdf(os.path.join(raw_X_data_dir, 'X_hf_order_quantile_data_v1.h5'),
                                          key='X_hf_order_quantile_data_v1',
                                          mode='w',
                                          format='table')

        print('get hf order quantile & daily trade & valuation & cap data done!\n')
    else:
        if not is_overwrite:
            hist_X_hf_order_quantile_data_v1 = \
                pd.read_hdf(os.path.join(raw_X_data_dir, 'X_hf_order_quantile_data_v1.h5'), start=-10)

            hist_end_date = hist_X_hf_order_quantile_data_v1.index.get_level_values('date').max()

            X_hf_order_quantile_data_v1 = \
                X_hf_order_quantile_data_v1[X_hf_order_quantile_data_v1.index.get_level_values('date') > hist_end_date]


            if X_hf_order_quantile_data_v1.shape[0] > 0:
                curr_start_date = X_hf_order_quantile_data_v1.index.get_level_values('date').min()

                X_hf_order_quantile_data_v1.to_hdf(os.path.join(raw_X_data_dir, 'X_hf_order_quantile_data_v1.h5'),
                                                   key='X_hf_order_quantile_data_v1',
                                                   mode='a',
                                                   append=True,
                                                   format='table')

                print('update hf order quantile & daily trade & valuation & cap data from %s to %s!\n' %
                      (curr_start_date.strftime('%Y-%m-%d'), end_date))
            else:
                print('hf order quantile & daily trade & valuation & cap data no update, since previous end date is %s!\n' %
                      (hist_end_date.strftime('%Y-%m-%d')))
        else:
            hist_X_hf_order_quantile_data_v1 = pd.read_hdf(os.path.join(raw_X_data_dir, 'X_hf_order_quantile_data_v1.h5'))

            curr_start_date = X_hf_order_quantile_data_v1.index.get_level_values('date').min()
            curr_end_date = X_hf_order_quantile_data_v1.index.get_level_values('date').max()

            hist_X_hf_order_quantile_data_v1 = \
                hist_X_hf_order_quantile_data_v1[hist_X_hf_order_quantile_data_v1.index.get_level_values('date') < curr_start_date]

            concat_X_hf_order_quantile_data_v1 = pd.concat([hist_X_hf_order_quantile_data_v1, X_hf_order_quantile_data_v1], axis=0)

            concat_X_hf_order_quantile_data_v1.to_hdf(os.path.join(raw_X_data_dir, 'X_hf_order_quantile_data_v1.h5'),
                                                 key='X_hf_order_quantile_data_v1',
                                                 mode='w',
                                                 format='table')

            print('update(overwrite) hf order quantile & daily trade & valuation & cap data from %s to %s!\n' %
                  (curr_start_date.strftime('%Y-%m-%d'),
                   curr_end_date.strftime('%Y-%m-%d')))


if __name__ == '__main__':
    # start_date = '2016-01-01'
    # end_date = '2020-01-31'
    # is_update = False
    # is_overwrite = True

    start_date = '2020-01-01'
    end_date = '2020-03-31'
    is_update = True
    is_overwrite = False

    main(start_date=start_date, end_date=end_date, is_update=is_update, is_overwrite=is_overwrite)

    # for test
    # X_hf_order_quantile_data_v1 = pd.read_hdf(os.path.join(raw_X_data_dir, 'X_hf_order_quantile_data_v1.h5'))
    # tmp = X_hf_order_quantile_data_v1[X_hf_order_quantile_data_v1.index.get_level_values('stock_code') == '300104.SZ']
    # tmp.loc[tmp.index.get_level_values('date') <= pd.to_datetime('2018-01-24')].dropna()